#!/bin/bash
"""服务启动脚本

用于启动项目所需的所有服务，包括Kafka、Flink、模型服务等。
支持单独启动某个服务或启动所有服务。
"""

# 配置参数
DOCKER_COMPOSE_FILE="${DOCKER_COMPOSE_FILE:-./docker-compose.yml}"
SERVICE_LIST=()
WAIT_TIMEOUT=${WAIT_TIMEOUT:-30}

# 日志函数
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# 检查Docker是否可用
check_docker() {
    if ! command -v docker &> /dev/null; then
        log "错误: 未找到Docker命令，请先安装Docker"
        return 1
    fi
    
    if ! docker info &> /dev/null; then
        log "错误: Docker服务未运行，请先启动Docker"
        return 1
    fi
    
    if ! command -v docker-compose &> /dev/null; then
        log "错误: 未找到docker-compose命令，请先安装docker-compose"
        return 1
    fi
    
    return 0
}

# 启动服务
start_service() {
    local service_name=$1
    
    log "开始启动服务: $service_name"
    
    case $service_name in
        "kafka")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d kafka
            wait_for_service "kafka" 9092
            ;;
        "zookeeper")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d zookeeper
            wait_for_service "zookeeper" 2181
            ;;
        "flink-jobmanager")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d flink-jobmanager
            wait_for_service "flink-jobmanager" 8081
            ;;
        "flink-taskmanager")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d flink-taskmanager
            # TaskManager没有直接的健康检查端口，可以检查日志
            ;;
        "redis")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d redis
            wait_for_service "redis" 6379
            ;;
        "model-service")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d model-service
            wait_for_service "model-service" 8000
            ;;
        "spider-service")
            docker-compose -f "$DOCKER_COMPOSE_FILE" up -d spider-service
            wait_for_service "spider-service" 8001
            ;;
        "all")
            log "开始启动所有服务..."
            # 按依赖顺序启动服务
            start_service "zookeeper"
            start_service "kafka"
            start_service "redis"
            start_service "flink-jobmanager"
            start_service "flink-taskmanager"
            start_service "model-service"
            start_service "spider-service"
            ;;
        *)
            log "错误: 不支持的服务: $service_name"
            return 1
            ;;
    esac
    
    if [ $? -eq 0 ]; then
        log "服务启动成功: $service_name"
        return 0
    else
        log "错误: 服务启动失败: $service_name"
        return 1
    fi
}

# 等待服务启动
wait_for_service() {
    local service_name=$1
    local port=$2
    local timeout=$WAIT_TIMEOUT
    local start_time=$(date +%s)
    
    log "等待服务 $service_name 在端口 $port 上启动..."
    
    while true; do
        # 检查容器是否运行
        if docker-compose -f "$DOCKER_COMPOSE_FILE" ps | grep "$service_name" | grep -q "Up"; then
            # 检查端口是否可访问
            if nc -z localhost "$port" 2>/dev/null; then
                log "服务 $service_name 已在端口 $port 上启动"
                return 0
            fi
        fi
        
        # 检查是否超时
        local current_time=$(date +%s)
        local elapsed=$((current_time - start_time))
        if [ $elapsed -gt $timeout ]; then
            log "错误: 服务 $service_name 启动超时（$timeout 秒）"
            return 1
        fi
        
        sleep 2
    done
}

# 停止服务
stop_service() {
    local service_name=$1
    
    log "开始停止服务: $service_name"
    
    if [ "$service_name" = "all" ]; then
        docker-compose -f "$DOCKER_COMPOSE_FILE" down
    else
        docker-compose -f "$DOCKER_COMPOSE_FILE" stop "$service_name"
    fi
    
    if [ $? -eq 0 ]; then
        log "服务停止成功: $service_name"
        return 0
    else
        log "错误: 服务停止失败: $service_name"
        return 1
    fi
}

# 显示服务状态
show_status() {
    log "显示所有服务状态..."
    docker-compose -f "$DOCKER_COMPOSE_FILE" ps
}

# 显示帮助信息
show_help() {
    echo "用法: $0 [命令] [服务名]"
    echo "服务启动脚本，用于启动项目所需的所有服务"
    echo ""
    echo "命令:"
    echo "  start                 启动服务"
    echo "  stop                  停止服务"
    echo "  status                显示服务状态"
    echo "  restart               重启服务"
    echo ""
    echo "服务名:"
    echo "  all                   所有服务"
    echo "  zookeeper             Zookeeper服务"
    echo "  kafka                 Kafka服务"
    echo "  redis                 Redis服务"
    echo "  flink-jobmanager      Flink JobManager"
    echo "  flink-taskmanager     Flink TaskManager"
    echo "  model-service         模型服务"
    echo "  spider-service        爬虫服务"
    echo ""
    echo "环境变量:"
    echo "  DOCKER_COMPOSE_FILE   Docker Compose文件路径"
    echo "  WAIT_TIMEOUT          服务启动等待超时时间（秒）"
}

# 主函数
main() {
    log "===== 服务管理脚本启动 ====="
    
    # 检查Docker是否可用
    if ! check_docker; then
        log "===== 服务管理脚本执行失败 ====="
        exit 1
    fi
    
    # 处理命令行参数
    if [ $# -eq 0 ]; then
        show_help
        log "===== 服务管理脚本执行失败 ====="
        exit 1
    fi
    
    local command=$1
    local service_name="all"
    
    if [ $# -ge 2 ]; then
        service_name=$2
    fi
    
    case $command in
        "start")
            start_service "$service_name"
            local result=$?
            ;;
        "stop")
            stop_service "$service_name"
            local result=$?
            ;;
        "status")
            show_status
            local result=0
            ;;
        "restart")
            stop_service "$service_name"
            start_service "$service_name"
            local result=$?
            ;;
        *)
            log "错误: 未知命令: $command"
            show_help
            local result=1
            ;;
    esac
    
    if [ $result -eq 0 ]; then
        log "===== 服务管理脚本执行成功 ====="
        exit 0
    else
        log "===== 服务管理脚本执行失败 ====="
        exit 1
    fi
}

# 执行主函数
main "$@"